﻿using System.Collections.Concurrent;
using System.Net.WebSockets;
using System.Text;

namespace XmppWebTest.Hub
{
    public class WebSocketManager
    {
        private readonly ConcurrentDictionary<string, WebSocket> _clients = new();

        public void AddClient(string clientId, WebSocket socket)
        {
            _clients[clientId] = socket;
        }

        public async Task RemoveClientAsync(string clientId)
        {
            if (_clients.TryRemove(clientId, out var socket))
            {
                await socket.CloseAsync(WebSocketCloseStatus.NormalClosure, "Connection closed", CancellationToken.None);
            }
        }

        public async Task SendMessageToAllAsync(string message)
        {
            var buffer = Encoding.UTF8.GetBytes(message);
            var tasks = _clients.Values
                .Where(socket => socket.State == WebSocketState.Open)
                .Select(socket => socket.SendAsync(new ArraySegment<byte>(buffer), WebSocketMessageType.Text, true, CancellationToken.None))
                .ToArray();

            await Task.WhenAll(tasks);
        }

        public async Task ReceiveMessagesAsync(string clientId, WebSocket webSocket)
        {
            var buffer = new byte[1024 * 4];

            while (webSocket.State == WebSocketState.Open)
            {
                var result = await webSocket.ReceiveAsync(new ArraySegment<byte>(buffer), CancellationToken.None);

                if (result.MessageType == WebSocketMessageType.Close)
                {
                    await RemoveClientAsync(clientId);
                    Console.WriteLine($"WebSocket 客户端 {clientId} 断开连接！");
                }
            }
        }
    }
}